################
#
# Deep Flow Prediction - N. Thuerey, K. Weissenov, H. Mehrotra, N. Mainali, L. Prantl, X. Hu (TUM)
#
# Dataset handling
#
################

from torch.utils.data import Dataset
import numpy as np
from os import listdir
import random

# global switch, use fixed max values for dim-less airfoil data?
fixedAirfoilNormalization = False
# global switch, make data dimensionless?
makeDimLess = True
# global switch, remove constant offsets from pressure channel?
removePOffset = True

## helper - compute absolute of inputs or targets
def find_absmax(data, use_targets, x):
    maxval = 0
    for i in range(data.totalLength):
        if use_targets == 0:
            temp_tensor = data.inputs[i]
        else:
            temp_tensor = data.targets[i]
        temp_max = np.max(np.abs(temp_tensor[x]))
        if temp_max > maxval:
            maxval = temp_max
    return maxval


######################################## DATA LOADER #########################################
#         also normalizes data with max , and optionally makes it dimensionless              #

def LoaderNormalizer(data, isTest=False, shuffle=0, dataProp=None):
    """
    # data: pass TurbDataset object with initialized dataDir / dataDirTest paths
    # train: when off, process as test data (first load regular for normalization if needed, then replace by test data)
    # dataProp: proportions for loading & mixing 3 different data directories "reg", "shear", "sup"
    #           should be array with [total-length, fraction-regular, fraction-superimposed, fraction-sheared],
    #           passing None means off, then loads from single directory
    """

    if dataProp is None:
        # load single directory
        files = listdir(data.dataDir)
        files.sort()
        # 打乱顺序
        for i in range(shuffle):
            random.shuffle(files) 
        #如果是测试模式，截断测试数据集，保留前十个
        if isTest:
            print("Reducing data to load for tests")
            files = files[0:min(10, len(files))]
        data.totalLength = len(files)

        # 初始分配4通道
        data.inputs  = np.empty((len(files), 4, 128, 128))  # 调整为4通道
        data.targets = np.empty((len(files), 3, 128, 128))

        for i, file in enumerate(files):
            npfile = np.load(data.dataDir + file)
            d = npfile['a']
            # 加载原始3通道数据
            data.inputs[i, :3] = d[0:3]  # 前3通道为原始输入
            data.targets[i] = d[3:6]
            # 添加管口掩码
            mask = np.zeros((1, 128, 128))
            # 假设管口位于输入数据的特定区域（需根据几何调整）
            mask[:, :, :10] = 1.0  # 示例：前10列为管口
            data.inputs[i, 3:4] = mask  # 第4通道为掩码
        print("Number of data loaded:", len(data.inputs))

    else:
        # load from folders reg, sup, and shear under the folder dataDir
        data.totalLength = int(dataProp[0])
        data.inputs  = np.empty((data.totalLength, 4, 128, 128))  # 调整为4通道
        data.targets = np.empty((data.totalLength, 3, 128, 128))

        files1 = listdir(data.dataDir + "reg/")
        files1.sort()
        files2 = listdir(data.dataDir + "sup/")
        files2.sort()
        files3 = listdir(data.dataDir + "shear/")
        files3.sort()
        
        for i in range(shuffle):
            random.shuffle(files1) 
            random.shuffle(files2) 
            random.shuffle(files3) 

        temp_1, temp_2 = 0, 0
        for i in range(data.totalLength):
            if i >= (1 - dataProp[3]) * dataProp[0]:
                npfile = np.load(data.dataDir + "shear/" + files3[i - temp_2])
                d = npfile['a']
                data.inputs[i, :3] = d[0:3]
                data.targets[i] = d[3:6]
            elif i >= (dataProp[1]) * dataProp[0]:
                npfile = np.load(data.dataDir + "sup/" + files2[i - temp_1])
                d = npfile['a']
                data.inputs[i, :3] = d[0:3]
                data.targets[i] = d[3:6]
                temp_2 = i + 1
            else:
                npfile = np.load(data.dataDir + "reg/" + files1[i])
                d = npfile['a']
                data.inputs[i, :3] = d[0:3]
                data.targets[i] = d[3:6]
                temp_1 = i + 1
                temp_2 = i + 1
            # 添加管口掩码
            mask = np.zeros((1, 128, 128))
            mask[:, :, :10] = 1.0  # 示例：前10列为管口
            data.inputs[i, 3:4] = mask
        print("Number of data loaded (reg, sup, shear):", temp_1, temp_2 - temp_1, i + 1 - temp_2)

    ################################## NORMALIZATION OF TRAINING DATA ##########################################

    if removePOffset:
        for i in range(data.totalLength):
            data.targets[i, 0, :, :] -= np.mean(data.targets[i, 0, :, :])
            data.targets[i, 0, :, :] -= data.targets[i, 0, :, :] * data.inputs[i, 2, :, :]  # 压力 * 掩码

    # make dimensionless based on current data set
    if makeDimLess:
        for i in range(data.totalLength):
            v_norm = (np.max(np.abs(data.inputs[i, 0, :, :]))**2 + np.max(np.abs(data.inputs[i, 1, :, :]))**2)**0.5
            data.targets[i, 0, :, :] /= v_norm**2
            data.targets[i, 1, :, :] /= v_norm
            data.targets[i, 2, :, :] /= v_norm

    # normalize to -1..1 range, from min/max of predefined
    if fixedAirfoilNormalization:
        data.max_inputs_0 = 100.
        data.max_inputs_1 = 38.12
        data.max_inputs_2 = 1.0
        data.max_inputs_3 = 1.0  # 管口掩码的最大值

        if makeDimLess:
            data.max_targets_0 = 4.65
            data.max_targets_1 = 2.04
            data.max_targets_2 = 2.37
            print("Using fixed maxima " + format([data.max_targets_0, data.max_targets_1, data.max_targets_2]))
        else:
            data.max_targets_0 = 40000.
            data.max_targets_1 = 200.
            data.max_targets_2 = 216.
            print("Using fixed maxima " + format([data.max_targets_0, data.max_targets_1, data.max_targets_2]))

    else:
        data.max_inputs_0 = find_absmax(data, 0, 0)
        data.max_inputs_1 = find_absmax(data, 0, 1)
        data.max_inputs_2 = find_absmax(data, 0, 2)
        data.max_inputs_3 = find_absmax(data, 0, 3)  # 管口掩码
        print("Maxima inputs " + format([data.max_inputs_0, data.max_inputs_1, data.max_inputs_2, data.max_inputs_3]))

        data.max_targets_0 = find_absmax(data, 1, 0)
        data.max_targets_1 = find_absmax(data, 1, 1)
        data.max_targets_2 = find_absmax(data, 1, 2)
        print("Maxima targets " + format([data.max_targets_0, data.max_targets_1, data.max_targets_2]))

    data.inputs[:, 0, :, :] *= (1.0 / data.max_inputs_0)
    data.inputs[:, 1, :, :] *= (1.0 / data.max_inputs_1)
    data.inputs[:, 2, :, :] *= (1.0 / data.max_inputs_2)
    data.inputs[:, 3, :, :] *= (1.0 / data.max_inputs_3)  # 归一化管口掩码

    data.targets[:, 0, :, :] *= (1.0 / data.max_targets_0)
    data.targets[:, 1, :, :] *= (1.0 / data.max_targets_1)
    data.targets[:, 2, :, :] *= (1.0 / data.max_targets_2)

    ###################################### NORMALIZATION OF TEST DATA #############################################

    if isTest:
        files = listdir(data.dataDirTest)
        files.sort()
        data.totalLength = len(files)
        data.inputs = np.empty((len(files), 4, 128, 128))  # 调整为4通道
        data.targets = np.empty((len(files), 3, 128, 128))
        for i, file in enumerate(files):
            npfile = np.load(data.dataDirTest + file)
            d = npfile['a']
            data.inputs[i, :3] = d[0:3]  # 前3通道为原始输入
            data.targets[i] = d[3:6]
            # 添加管口掩码
            mask = np.zeros((1, 128, 128))
            mask[:, :, :10] = 1.0  # 示例：前10列为管口
            data.inputs[i, 3:4] = mask

        if removePOffset:
            for i in range(data.totalLength):
                data.targets[i, 0, :, :] -= np.mean(data.targets[i, 0, :, :])
                data.targets[i, 0, :, :] -= data.targets[i, 0, :, :] * data.inputs[i, 2, :, :]

        if makeDimLess:
            for i in range(len(files)):
                v_norm = (np.max(np.abs(data.inputs[i, 0, :, :]))**2 + np.max(np.abs(data.inputs[i, 1, :, :]))**2)**0.5
                data.targets[i, 0, :, :] /= v_norm**2
                data.targets[i, 1, :, :] /= v_norm
                data.targets[i, 2, :, :] /= v_norm

        data.inputs[:, 0, :, :] *= (1.0 / data.max_inputs_0)
        data.inputs[:, 1, :, :] *= (1.0 / data.max_inputs_1)
        data.inputs[:, 2, :, :] *= (1.0 / data.max_inputs_2)
        data.inputs[:, 3, :, :] *= (1.0 / data.max_inputs_3)

        data.targets[:, 0, :, :] *= (1.0 / data.max_targets_0)
        data.targets[:, 1, :, :] *= (1.0 / data.max_targets_1)
        data.targets[:, 2, :, :] *= (1.0 / data.max_targets_2)

    print("Data stats, input mean %f, max %f; targets mean %f , max %f " % (
        np.mean(np.abs(data.targets), keepdims=False), np.max(np.abs(data.targets), keepdims=False),
        np.mean(np.abs(data.inputs), keepdims=False), np.max(np.abs(data.inputs), keepdims=False)))

    return data

######################################## DATA SET CLASS #########################################

class TurbDataset(Dataset):
    TRAIN = 0
    TEST = 2

    def __init__(self, dataProp=None, mode=TRAIN, dataDir="../data/train/", dataDirTest="../data/test/", shuffle=0, normMode=0):
        global makeDimLess, removePOffset

        if not (mode == self.TRAIN or mode == self.TEST):
            print("Error - TurbDataset invalid mode " + format(mode))
            exit(1)

        if normMode == 1:    
            print("Warning - poff off!!")
            removePOffset = False
        if normMode == 2:    
            print("Warning - poff and dimless off!!!")
            makeDimLess = False
            removePOffset = False

        self.mode = mode
        self.dataDir = dataDir
        self.dataDirTest = dataDirTest

        # 加载并归一化数据
        self = LoaderNormalizer(self, isTest=(mode == self.TEST), dataProp=dataProp, shuffle=shuffle)

        if not self.mode == self.TEST:
            targetLength = self.totalLength
            self.inputs = self.inputs[:targetLength]
            self.targets = self.targets[:targetLength]
            self.totalLength = self.inputs.shape[0]

    def __len__(self):
        return self.totalLength

    def __getitem__(self, idx):
        # 返回输入、目标和管口掩码（从输入的第4通道提取）
        return self.inputs[idx], self.targets[idx], self.inputs[idx, 3:4]

    def denormalize(self, data, v_norm):
        a = data.copy()
        a[0, :, :] /= (1.0 / self.max_targets_0)
        a[1, :, :] /= (1.0 / self.max_targets_1)
        a[2, :, :] /= (1.0 / self.max_targets_2)

        if makeDimLess:
            a[0, :, :] *= v_norm**2
            a[1, :, :] *= v_norm
            a[2, :, :] *= v_norm
        return a



